{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def step(self, act):\n",
    "    nLevel = 3\n",
    "\n",
    "    n_chosen_act0 = sum(act==0)\n",
    "    n_chosen_act1 = sum(act==1)\n",
    "    if self.TS==0:  # Level 1\n",
    "        # we expect all but one agents choose act1, only one agent choose act0\n",
    "        # reward: \n",
    "        #   - Any agent choose act0: Team Reward +1\n",
    "        #   - No agent choose act0: Team Reward +0\n",
    "        reward = n_chosen_act0/self.n_agents\n",
    "        self.TS = 1\n",
    "        ob = self.get_obs(TS=1)\n",
    "\n",
    "    elif self.TS==1:  # Level 2\n",
    "        # we expect all but one agents choose act1, only one agent choose act0\n",
    "        # reward: \n",
    "        #   - Any agent choose act0: Team Reward +1\n",
    "        #   - No agent choose act0: Team Reward +0\n",
    "        reward = -1 if n_chosen_act0==0 else +1\n",
    "        self.TS = 2\n",
    "        ob = self.get_obs(TS=2)\n",
    "\n",
    "    elif self.TS==2:  # Level 3\n",
    "        # we expect all but one agents choose act0, only one agent choose act1\n",
    "        # reward: \n",
    "        #   - Any agent choose act1: Team Reward +1\n",
    "        #   - No agent choose act1: Team Reward +0\n",
    "        reward = -1 if n_chosen_act1==0 else +1\n",
    "        self.TS = 3\n",
    "        ob = self.get_obs(TS=3) # Terminal obs, won't be accepted\n",
    "    else:\n",
    "        assert False, 'Should not be here !'\n",
    "\n",
    "    # obs: a Tensor with shape (n_agent, ...)\n",
    "    if self.TS >= nLevel: \n",
    "        done = True\n",
    "    else: \n",
    "        done = False\n",
    "    # \n",
    "    info = {'win':False}\n",
    "    reward_allteam = np.array([reward])\n",
    "    return (ob, reward_allteam, done, info)\n",
    "\n",
    "\n",
    "'''\n",
    "\n",
    "    if self.TS==0:\n",
    "        # we expect all but one agents choose act1, only one agent choose act0\n",
    "        # reward: \n",
    "        #   - Any agent choose act0: Team Reward +1\n",
    "        #   - No agent choose act0: Team Reward +0\n",
    "        reward = n_chosen_act0/self.n_agents\n",
    "\n",
    "    Results:\n",
    "          act0    act1\n",
    "        [[9.9981e-01, 1.9052e-04],\n",
    "         [9.9962e-01, 3.7964e-04],\n",
    "         [9.9970e-01, 2.9723e-04],\n",
    "         [9.9981e-01, 1.9268e-04],\n",
    "         [9.9986e-01, 1.4075e-04]],\n",
    "     _____________________________________________________________________________\n",
    "\n",
    "        elif self.TS==1:  # Level 2\n",
    "            # we expect all but one agents choose act1, only one agent choose act0\n",
    "            # reward: \n",
    "            #   - Any agent choose act0: Team Reward +1\n",
    "            #   - No agent choose act0: Team Reward +0\n",
    "            reward = -1 if n_chosen_act0==0 else +1\n",
    "            self.TS = 2\n",
    "            ob = self.get_obs(TS=2)\n",
    "        [[0.7883, 0.2117],\n",
    "         [0.5876, 0.4124],\n",
    "         [0.7635, 0.2365],\n",
    "         [0.8501, 0.1499],\n",
    "         [0.8389, 0.1611]],\n",
    "\n",
    "     _____________________________________________________________________________\n",
    "\n",
    "        elif self.TS==2:  # Level 3\n",
    "            # we expect all but one agents choose act0, only one agent choose act1\n",
    "            # reward: \n",
    "            #   - Any agent choose act1: Team Reward +1\n",
    "            #   - No agent choose act1: Team Reward +0\n",
    "            reward = -1 if n_chosen_act1==0 else +1\n",
    "            self.TS = 3\n",
    "\n",
    "          act0    act1\n",
    "        [[0.3068, 0.6932],\n",
    "         [0.3439, 0.6561],\n",
    "         [0.3616, 0.6384],\n",
    "         [0.3519, 0.6481],\n",
    "         [0.3549, 0.6451]],\n",
    "'''"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
